# 将3D的轨迹转换为2D（KITTI格式）

import numpy as np
import matplotlib.pyplot as plt
import os

import trans_func
import draw_traj

"""
将3D的轨迹转换为2D（KITTI格式）
"""


def kitti_traj_3d_to_2d(filepath, save_path):
    dataset = np.loadtxt(filepath, delimiter=' ')
    for i in range(len(dataset)):
        list_line = dataset[i]  # 读一行

        list_line.resize(3, 4)  # 向量转矩阵

        # 读取t, R数据
        t = np.array(list_line[:, 3])
        R = np.array(list_line[:, :3])

        # 转2D
        t[1] = 0
        euler = trans_func.rot2euler(R)
        euler[0] = 0
        euler[2] = 0
        R_2d = trans_func.euler2rot(euler)
        list_line[:, :3] = R_2d
        list_line[:, 3] = t

        list_line.resize(12)
        dataset[i] = list_line
    # 保存
    np.savetxt(save_path, dataset, delimiter=' ')


"""
将3D的轨迹转换为2D（Urban格式）
"""


def traj_3d_to_2d(filepath, save_path, datasetname="urban"):
    dataset = np.loadtxt(filepath, delimiter=' ')
    for i in range(len(dataset)):
        list_line = dataset[i]  # 读一行

        # 读取time, t, R数据
        # time = np.array(list_line[0])
        t = np.array(list_line[1:4])
        q = np.array(list_line[4:])

        if datasetname == "kitti":
            # 转2D
            t[1] = 0
            euler = trans_func.quaternion2euler(q)
            euler[0] = 0
            euler[2] = 0
        elif datasetname == "urban":
            # 转2D
            t[2] = 0
            euler = trans_func.quaternion2euler(q)
            euler[0] = 0
            euler[1] = 0

        q_2d = trans_func.euler2quaternion(euler)
        list_line[1:4] = t
        list_line[4:] = q_2d

        dataset[i] = list_line
    # 保存
    np.savetxt(save_path, dataset, delimiter=' ')


if __name__ == '__main__':
    # # 这里需要更改数据位置
    # # file_path = "/media/daybeha/LTFM2/dataset/kitti/sequences/poses"
    # # traj_3d = "08.txt"
    # # traj_2d = "08_2d.txt"
    #
    file_path = "//home/daybeha/Documents/github/Hybird-NeuroSLAM/vins_ws/results/"
    # traj_3d = "trajectory_kitti.txt"
    # traj_2d = "trajectory_kitti_2d.txt"
    traj_3d = "VIO_urban39_keyframe.txt"
    traj_2d = "VIO_urban39_keyframe_2d.txt"
    # #
    # file_path = "/home/daybeha/Documents/github/Hybird-NeuroSLAM/ORB_SLAM3_detailed_comments-master/Examples/src/results/"
    # # traj_3d = "trajectory_kitti.txt"
    # # traj_2d = "trajectory_kitti_2d.txt"
    # traj_3d = "stereo_kitti08.txt"
    # traj_2d = "stereo_kitti08_2d.txt"

    # kitti_traj_3d_to_2d(os.path.join(file_path, traj_3d), os.path.join(file_path, traj_2d))
    # draw_traj.draw_kitti(os.path.join(file_path, traj_2d))






    # file_path = "/media/daybeha/LTFM2/dataset/kitti/sequences/poses"
    # traj_3d = "08_tum.txt"
    # traj_2d = "08_tum_2d.txt"


    # # file_path = "/home/daybeha/Documents/github/dso_ws/src/results/"
    # file_path = "/home/daybeha/Documents/github/Hybird-NeuroSLAM/vins_ws/src/VINS-Fusion/results/"
    # # # file_path = "/home/daybeha/Documents/github/catkin_ws/src/results/"
    # traj_3d = "trajectory_tum.txt"
    # traj_2d = "trajectory_tum_2d.txt"

    # file_path = "/home/daybeha/Documents/github/catkin_ws/"
    # traj_3d = "vio.txt"
    # traj_2d = "vio_2d.txt"


    # file_path = "/media/daybeha/LTFM2/dataset/urban/urban30-gangnam/"
    # traj_3d = "global_pose_tum.txt"
    # traj_2d = "global_pose_tum_2d.txt"


    # file_path = "/home/daybeha/Documents/github/openvins_ws/src/open_vins/results/"
    # file_path = "/home/daybeha/Documents/github/myratslam_ws/src/ratslam_ros/results/fromopenvins_without_vm/Fast/urban28/"
    # traj_3d = "traj_estimate_tum.txt"
    # traj_2d = "traj_estimate_tum_2d.txt"

    traj_3d_to_2d(os.path.join(file_path, traj_3d), os.path.join(file_path, traj_2d), "urban")
    draw_traj.draw_tum(os.path.join(file_path, traj_2d), "urban")

    # file_path = "/home/daybeha/Documents/github/myratslam_ws/src/ratslam_ros/results/urban/urban39/vins_opt.txt"
    # # urban_traj_3d_to_2d(file_path, os.path.join(file_path, traj_2d), "urban")
    # draw_traj.draw_tum(file_path, "urban")
